#ifndef _SCAN_KERNEL_H_
#define _SCAN_KERNEL_H_

#include <cutil_inline.h>
#include <cutil_math.h>
#include <stdio.h>
#include "sat.h"
#include "scan.h"

// performs inclusive scan
template <typename T, ScanKind kind>
__global__ void kernel_scan(T *dev_src, T *dev_dst, unsigned int len, unsigned int stride)
{
	unsigned int k = blockIdx.x*blockDim.x + threadIdx.x;
	unsigned int idx = threadIdx.x;
	__shared__ T src[2*SCAN_BLOCK_DIM];
	__shared__ T dst[SCAN_BLOCK_DIM];

	src[2*idx] = (kind == Forward) ? dev_src[k] : dev_src[len - k - 1];
	__syncthreads(); // for coalescing

	if (k >= stride)
		src[2*idx+1] = (kind == Forward) ? dev_src[k-stride] : dev_src[len - k - 1 + stride];
	__syncthreads(); // prevent operating on nonexistent data

	if (k < len) {
		//dev_dst[k] = (k >= stride) ? dev_src[k - stride] + dev_src[k] : dev_src[k];
		dst[(kind == Forward) ? idx : min(SCAN_BLOCK_DIM, len-blockIdx.x*blockDim.x) - idx - 1] = (k >= stride) ? src[2*idx+1] + src[2*idx] : src[2*idx];
	}
	__syncthreads();
	if (k < len) {
		unsigned int dev_idx = (kind == Forward) ? k :
			((blockIdx.x == gridDim.x-1) ? 0 : len % SCAN_BLOCK_DIM + (gridDim.x - blockIdx.x - 1 - (len % SCAN_BLOCK_DIM != 0))*SCAN_BLOCK_DIM) + idx;
		//dev_dst[(kind == Forward) ? k : (gridDim.x - blockIdx.x - 1)*blockDim.x + idx] = dst[idx];
		dev_dst[dev_idx] = dst[idx];
	}
}

template <typename T>
__global__ void kernel_sat_scan(T *dev_src, T * dev_dst, unsigned int len, unsigned int stride)
{
	unsigned int k = blockIdx.x*blockDim.x + threadIdx.x;
	unsigned int idx = threadIdx.x;
	unsigned int offset = len * blockIdx.y;
	__shared__ T src[2*SCAN_BLOCK_DIM];

	src[2*idx] = dev_src[k + offset];
	__syncthreads(); // for coalescing

	if (k >= stride)
		src[2*idx+1] = dev_src[k-stride + offset];
	__syncthreads(); // prevent operating on nonexistent data

	if (k < len) {
		dev_dst[k + offset] = (k >= stride) ? src[2*idx+1] + src[2*idx] : src[2*idx];
	}
}



// got a clever implementation of this courtesy of: http://www.gamedev.net/topic/229831-nearest-power-of-2/
unsigned NextPow2( unsigned n )
{
	const int MantissaMask = (1<<23) - 1;

	(*(float*)&n) = (float) n;
	n = n + MantissaMask & ~MantissaMask;
	n = (unsigned) (*(float*)&n);

	return n;
} 

template <typename T, ScanKind kind>
void scan_helper(T *dev_src, T *temp, unsigned int len, unsigned int stride, unsigned int numBlocks, unsigned int blockSize)
{
	dim3 block(blockSize, 1, 1);
	dim3 grid(numBlocks, 1, 1);


#ifdef DEBUG
	T arr2[4096];
	cutilSafeCall(cudaMemcpy(&arr2, dev_src, sizeof(T) * len, cudaMemcpyDeviceToHost));
	printf("Stride: %u\n", stride);
	for (unsigned int i = 0; i < len; i++) {
		printf("%3.0f, ", arr2[i].x);
	}
	printf("\n");
#endif

	kernel_scan<T, kind><<<grid, block>>>(dev_src, temp, len, stride);

#ifdef DEBUG
	T arr[4096];
	cutilSafeCall(cudaMemcpy(&arr, temp, sizeof(T) * len, cudaMemcpyDeviceToHost));
	for (unsigned int i = 0; i < len; i++) {
		printf("%3.0f, ", arr[i].x);
	}
	printf("\n");
#endif
	
	// need an additional iteration of not a power of 2
	if (stride == NextPow2((unsigned int)ceil((float)len / 2))) {
		cutilSafeCall(cudaMemcpy(dev_src, temp, len * sizeof(T), cudaMemcpyDeviceToDevice));
	}
	else {
		// ping pong buffers
		T * swap = dev_src;
		dev_src = temp;
		temp = swap;
		scan_helper<T, kind>(dev_src, temp, len, stride << 1, numBlocks, blockSize);
	}
}


// SAT scan helper
template <typename T>
T * sat_scan_helper(T *dev_src, T *dev_dst, const unsigned int width, const unsigned int height, 
					 unsigned int stride, unsigned int numBlocks, unsigned int blockSize)
{
	dim3 block(blockSize, 1, 1);
	dim3 grid(numBlocks, height, 1);

#ifdef DEBUG
	T arr2[16];
	cutilSafeCall(cudaMemcpy(&arr2, dev_src+4351, sizeof(T) * 16, cudaMemcpyDeviceToHost));
	printf("Stride: %u\n", stride);
	for (unsigned int i = 0; i < 16; i++) {
		printf("%3.0f, ", arr2[i].x);
	}
	printf("\n");
#endif

	kernel_sat_scan<T><<<grid, block>>>(dev_src, dev_dst, width, stride);
	//cudaThreadSynchronize();

#ifdef DEBUG
	T arr[16];
	cutilSafeCall(cudaMemcpy(&arr, dev_dst+4351, sizeof(T) * 16, cudaMemcpyDeviceToHost));
	for (unsigned int i = 0; i < 16; i++) {
		printf("%3.0f, ", arr[i].x);
	}
	printf("\n");
#endif

	// need an additional iteration of not a power of 2
	if (stride == NextPow2((unsigned int)ceil((float)width / 2))) {
		return dev_dst;
	}
	else {
		// ping pong buffers
		T * temp = dev_src;
		dev_src = dev_dst;
		dev_dst = temp;
		return sat_scan_helper<T>(dev_src, dev_dst, width, height, stride << 1, numBlocks, blockSize);
	}
}

template float4 * sat_scan<float4>(SAT_test_t *t, const unsigned int width, const unsigned int height);

template <typename T, ScanKind kind>
void scan(T* dev_src, const unsigned int len) {
	cudaDeviceProp devProp;
	cudaGetDeviceProperties(&devProp, 0);

	unsigned int blockSize = SCAN_BLOCK_DIM;
	unsigned int numBlocks = (unsigned int) ceil( (float) len / blockSize );

	if (len * sizeof(T) >  devProp.totalGlobalMem) {
		fprintf(stderr, "Too big, failing.\n");
		return;
	}

	T *temp_d;
	cutilSafeCall(cudaMalloc(&temp_d, len * sizeof(T)));
	// make sure first value is initted since algo won't touch it
	cutilSafeCall(cudaMemcpy(temp_d, dev_src, sizeof(T), cudaMemcpyDeviceToDevice));

	scan_helper<T, kind>(dev_src, temp_d, len, 1, numBlocks, blockSize);
	cutilSafeCall(cudaFree(temp_d));
}


template <typename T>
T * sat_scan(SAT_test_t *t, const unsigned int width, const unsigned int height) {
	cudaDeviceProp devProp;
	cudaGetDeviceProperties(&devProp, 0);

	unsigned int blockSize = t->blockSize;
	unsigned int numBlocks = (unsigned int) ceil( (float) width / blockSize );

	if (width * height * sizeof(T) > devProp.totalGlobalMem) {
		fprintf(stderr, "Too big, failing.\n");
		return NULL;
	}

	return sat_scan_helper<T>(t->dev_src, t->dev_dst, width, height, 1, numBlocks, blockSize);
}

#endif // #ifndef _SCAN_KERNEL_H_
